# ================================================================================== #
# Complaints About Police Misconduct Have Adverse Effects for Black Civilians (PSRM)
# - Script: Load required packages and custom functions
# - Author: Patrick Kraft (patrickwilli.kraft@uc3m.es)
# ================================================================================== #


# Required packages -------------------------------------------------------

library(here)
library(tseries)
library(forecast)
library(lmtest)
library(sandwich)
library(vars)
library(gridExtra)
library(tidyr)
library(tidyverse)
library(xtable)
library(Amelia)


# Set seed for reproducibility --------------------------------------------

set.seed(20150503)


# Plot defaults -----------------------------------------------------------

plot_default <- theme_classic(base_size = 8) + 
  theme(panel.border = element_rect(fill=NA))


# Custom functions --------------------------------------------------------

prePlot <- function(mlist, vars = NULL, report = NULL, NeweyWest = TRUE){
  ## function to prepare results from lm result list for plotting with ggplot
  ## mlist: list of lm model results (can also be from grangerlm function)
  ## vars = vector of variable names to be used
  ## report = vector of variable names to be included in the output (matches via grep)
  ## NeweyWest = report Newey West standard errors (caution, does not work with VARs)
  ## output = table of coefficients, SE (Newey West), Variable name, model name
  out <- NULL
  
  extractLm <- function(mlist, vars, report, NeweyWest){
    res <- NULL
    
    if(class(mlist) == "varest"){
      if(NeweyWest == T) stop("Newey West SEs not compatible with VAR result")
      mlist <- mlist$varresult
      varest <- TRUE
    }
    
    for(i in 1:length(mlist)){
      if(NeweyWest) {
        tmp <- coeftest(mlist[[i]], vcov=NeweyWest(mlist[[i]]))
      } else tmp <- coeftest(mlist[[i]])
      tmp <- data.frame(Estimate = tmp[,"Estimate"], SE = tmp[,"Std. Error"])
      if(!is.null(report)) tmp <- tmp[as.vector(sapply(report, grep, rownames(tmp))),]
      if(!is.null(vars)) {
        tmp$Variables <- vars
      } else {
        tmp$Variables <- rownames(tmp)
      }
      tmp$Model <- names(mlist)[i]
      res <- rbind(res, tmp)
    }
    res
  }
  
  if(class(mlist[[1]]) != "list" & class(mlist[[1]]) != "varest"){
    out <- extractLm(mlist, vars = vars, report = report, NeweyWest = NeweyWest)
  } else {
    for(j in 1:length(mlist)){
      tmp <- extractLm(mlist[[j]], vars = vars, report = report, NeweyWest = NeweyWest)
      tmp$Group <- names(mlist)[j]
      tmp$Variables <- gsub(paste0(".",names(mlist)[j]),"", tmp$Variables)
      tmp$Model <- gsub(paste0(".",names(mlist)[j]),"", tmp$Model)
      out <- rbind(out, tmp)
    }
  }
  
  rownames(out) <- NULL
  return(out)
}


extractCoefs <- function(x, model = "sqf_d"){
  filter(prePlot(x$varresult, NeweyWest = F, report = "ccrb"), Model == model)
}


varSim <- function (object, ..., n.ahead = 6, ci = 0.95, nsim=1000, 
                    dumvar, shock, level, vars = 3){
  ###########################################################
  # predict series based on VAR result
  # parts of the code adapted from vars:::predict.varest
  ###########################################################
  
  ## extract model objects
  K <- object$K                  # dimension of the VAR (# endogenous)
  p <- object$p                  # lag order of the VAR
  obs <- object$obs              # no. observations
  type <- object$type            # type
  data.all <- object$datamat     # full data.frame
  ynames <- colnames(object$y)   # names of endogenous variables
  xnames <- c("const", rownames(coef(object)[[1]])
              [rownames(coef(object)[[1]])!="const"]
  )                  # varnames in same order as vcov
  n.ahead <- as.integer(n.ahead) # number of time steps
  Z <- object$datamat[, -c(1:K)] # data frame of predictors
  levelm <- matrix(rep(level, each=n.ahead)
                   , ncol=length(level)) # level as matrix
  
  
  ## simulate parameters
  Bcoef <- unlist(lapply(coef(object), function(x) x[xnames,1]))
  Bvcov <- vcov(object)
  Bsim <- mvrnorm(n = nsim, mu = Bcoef, Sigma = Bvcov)
  
  ## object to store results
  res <- NULL
  
  ## compute predicted values
  for(i in 1:nsim){
    ## arrange parameters in matrix for different series
    B <- matrix(Bsim[i,], nrow=vars, byrow = T)
    rownames(B) <- ynames
    colnames(B) <- xnames
    
    ## adjust order of parameters (endogenous, constant, exogenous)
    B <- B[,c(grep("\\.l\\d$",colnames(B)),grep("\\.l\\d$",colnames(B),invert=T))]
    
    ## matrix of exogenous vars + constant
    Zdet <- as.matrix(data.frame(const=1, dumvar))
    
    ## start with constant endogenous variables (change?)
    lasty <- rep(0, (K * (p + 1)))
    
    ## forecast
    forecast <- matrix(NA, ncol = K, nrow = n.ahead)
    colnames(forecast) <- ynames
    for (j in 1:(n.ahead)) {
      lasty <- lasty[1:(K * p)]
      Z <- c(lasty, Zdet[j, ])
      if(j == 1) {          ## first iteration stays constant
        forecast[j, ] <- 0
      } else if(j == 2) {          ## second iteration only introduces mean
        forecast[j, ] <- 0
        forecast[j, names(shock)] <- unlist(shock)
      } else {              ## every following iteration based on simulation
        forecast[j, ] <- B %*% Z
      }
      temp <- forecast[j, ]
      lasty <- c(temp, lasty)
    }
    
    ## generate level variable (all shocks are at the same level)
    
    
    res <- rbind(res, data.frame(forecast, apply(forecast,2,cumsum)+levelm
                                 , time = -1:(n.ahead-2), iter = i))
  }
  return(res)
}


#' Simulate expected values/first differences (replaces Zelig)
#' built: 2016-08-27, Patrick Kraft
#' @importFrom MASS mvrnorm
#' @importFrom sandwich vcovHC
#' @param models: list of model results (lm, glm, or vglm/tobit)
#' @param iv: data frame containing the values for comparison (only 2 rows, selected variables)
#' @param se: NULL, robust, newey; types of standard errors that should be used
#' @param nsim: number of simulations
#' @return data.frame: contains expected values, confidence intervals, variable names
#' @export
#'
sim <- function(models, iv, se=NULL, ci=c(0.025,0.975), nsim = 1000){
  
  ## prepare output object, convert input to model list
  out <- NULL
  if(class(models)[1] != "list") models <- list(models)
  
  for(i in 1:length(models)){
    ## simulate betas from sampling distribution
    if(is.null(se)){
      betas <- MASS::mvrnorm(nsim, coef(models[[i]]), sandwich::vcovHC(models[[i]]))
    } else if(se == "robust"){
      betas <- MASS::mvrnorm(nsim, coef(models[[i]]), vcov(models[[i]]))
    } else if(se == "newey"){
      betas <- MASS::mvrnorm(nsim, coef(models[[i]]), sandwich::NeweyWest(models[[i]]))
    }
    
    ## extract variable names
    vars <- names(coef(models[[i]]))
    int <- grep("[^)]:", vars)
    varsInt <- strsplit(vars[int], ":")
    
    ## generate matrix of covariates
    X <- matrix(1, nrow=length(vars), ncol=nrow(iv))
    X[vars %in% names(iv),] <- t(iv[vars[vars %in% names(iv)]])
    if(class(models[[i]])[1]=="lm"){
      means <- apply(models[[i]]$model[vars[-c(1,which(vars %in% names(iv)),int)]]
                     , 2, mean, na.rm=T)
    } else if(class(models[[i]])[1] == "glm"){
      means <- apply(models[[i]]$data[vars[-c(1,which(vars %in% names(iv)),int)]]
                     , 2, mean, na.rm=T)
    } else if(class(models[[i]])[1] == "vglm" & models[[i]]@family@vfamily == "tobit"){
      means <- apply(models[[i]]@x[,vars[-c(1,2,which(vars %in% names(iv)),int)]]
                     , 2, mean, na.rm=T)
    } else stop("Model type not supported")
    X[vars %in% names(means),] <- means
    
    ## calculate interaction effects
    if(length(varsInt)>0){
      for(j in 1:length(varsInt)){
        X[int[j],] <- apply(X[vars %in% varsInt[[j]],],2,prod)
      }
    }
    
    ## calculate expected values
    if(class(models[[i]])[1]=="lm"){
      evs <- betas %*% X
    } else if(class(models[[i]])[1] == "glm"){
      if(models[[i]]$family$link == "logit"){
        evs <- 1/(1+exp(-betas %*% X))
      } else if(models[[i]]$family$link == "probit"){
        evs <- pnorm(betas %*% X)
      } else stop("Model type not supported")
    } else if(class(models[[i]])[1] == "vglm" & models[[i]]@family@vfamily == "tobit"){
      ## IDEA: decompose effect of tobit in dP(Y>0) and dY|Y>0
      ## based on predicted values (rather than EVs)
      ## note that betas[,2] is log(Sigma) estimate
      ## CHECK CALCULATIONS!
      if(unique(models[[i]]@misc$Upper)!=Inf) stop("Upper limit not supported")
      if(unique(models[[i]]@misc$Lower)!=0) warning("Limit != 0 not testes yet!")
      loLim <- unique(models[[i]]@misc$Lower)[1,1]
      
      ## expected values for z>0
      evsTemp <- betas[,-2] %*% X[-2,]
      evs <- evsTemp + exp(betas[,2]) * dnorm(evsTemp/exp(betas[,2])) / pnorm(evsTemp/exp(betas[,2]))
      
      ## probability of z>0
      pvs <- array(dim = c(nsim,ncol(X),nsim))
      for(j in 1:nrow(pvs)){
        pvs[j,,] <- matrix(rnorm(nsim*ncol(X), mean = evsTemp[j,], sd = exp(betas[j,2]))
                           , ncol = nsim)
      }
      prob <- apply(pvs, 2, function(x) apply(x, 1, function(x) mean(x>loLim)))
    } else stop("Model type not supported")
    
    skip <- F
    
    if(nrow(iv)==2){
      ## calculate first differences
      evs <- evs[,2] - evs[,1]
      if(class(models[[i]])[1] == "vglm"){
        if(models[[i]]@family@vfamily == "tobit")
          prob <- prob[,2] - prob[,1]
      }
    } else if(nrow(iv)==4) {
      ## calculate difference-in-difference
      evs <- (evs[,2] - evs[,1]) - (evs[,4] - evs[,3])
      if(class(models[[i]])[1] == "vglm"){
        if(models[[i]]@family@vfamily == "tobit")
          prob <- (prob[,2] - prob[,1]) - (prob[,4] - prob[,3])
      }
    } else {
      ## compute predicted values for each step
      warning("Check number of scenarios - STILL TESTING")
      if(class(models[[i]])[1] != "vglm"){
        res <- data.frame(mean = apply(evs, 2, mean)
                          , cilo = apply(evs, 2, quantile, ci[1])
                          , cihi = apply(evs, 2, quantile, ci[2])
                          , dv = as.factor(colnames(models[[i]]$model)[1])
                          , iv = as.factor(paste(colnames(iv), collapse = "_"))
                          , ivval = iv[,1])
      } else if(models[[i]]@family@vfamily == "tobit"){
        res <- data.frame(mean = c(apply(prob, 2, mean), apply(evs, 2, mean))
                          , cilo = c(apply(prob, 2, quantile, ci[1]), apply(evs, 2, quantile, ci[1]))
                          , cihi = c(apply(prob, 2, quantile, ci[2]), apply(evs, 2, quantile, ci[2]))
                          , dv = as.factor(sub("(.*) \\~.*", "\\1", models[[i]]@call[2]))
                          , iv = as.factor(paste(colnames(iv), collapse = "_"))
                          , ivval = iv[,1]
                          , value = factor(rep(c("Probability P(y>0)","Expected Value E(y|y>0)"), each = nrow(iv))
                                           , levels = c("Probability P(y>0)","Expected Value E(y|y>0)")))
      } else stop("Check model type")
      out <- rbind(out, res)
      skip <- T
    }
    
    ## warning for Inf/-Inf in single iterations
    if(Inf %in% evs|-Inf %in% evs){
      warning(paste0("Inf/-Inf in ",length(evs[evs==Inf])+length(evs[evs==-Inf])," evs iteration(s)"))
      evs[evs==Inf|evs==-Inf] <- NA
    }
    
    if(!skip){
      ## generate output table
      if(class(models[[i]])[1] != "vglm"){
        res <- data.frame(mean = mean(evs)
                          , cilo = quantile(evs, ci[1])
                          , cihi = quantile(evs, ci[2])
                          , dv = as.factor(colnames(models[[i]]$model)[1])
                          , iv = as.factor(paste(colnames(iv), collapse = "_")))
      } else {
        res <- data.frame(mean = c(mean(prob, na.rm = T), mean(evs, na.rm = T))
                          , cilo = c(quantile(prob, ci[1], na.rm = T),quantile(evs, ci[1], na.rm = T))
                          , cihi = c(quantile(prob, ci[2], na.rm = T), quantile(evs, ci[2], na.rm = T))
                          , dv = as.factor(sub("(.*) \\~.*", "\\1", models[[i]]@call[2]))
                          , iv = as.factor(paste(colnames(iv), collapse = "_"))
                          , value = factor(c("Probability P(y>0)","Expected Value E(y|y>0)")
                                           , levels = c("Probability P(y>0)","Expected Value E(y|y>0)")))
      }
      out <- rbind(out, res)
    }
  }
  
  ## return output table
  rownames(out) <- NULL
  return(out)
}


adf.test2 <- function (x, alternative = c("stationary", "explosive")
                       , k = trunc((length(x) - 1)^(1/3))) 
{
  ## modification of output in adf.test (include "<" and ">" in p-value output rather than warning)
  if ((NCOL(x) > 1) || is.data.frame(x)) 
    stop("x is not a vector or univariate time series")
  if (any(is.na(x))) 
    stop("NAs in x")
  if (k < 0) 
    stop("k negative")
  alternative <- match.arg(alternative)
  DNAME <- deparse(substitute(x))
  k <- k + 1
  x <- as.vector(x, mode = "double")
  y <- diff(x)
  n <- length(y)
  z <- embed(y, k)
  yt <- z[, 1]
  xt1 <- x[k:n]
  tt <- k:n
  if (k > 1) {
    yt1 <- z[, 2:k]
    res <- lm(yt ~ xt1 + 1 + tt + yt1)
  }
  else res <- lm(yt ~ xt1 + 1 + tt)
  res.sum <- summary(res)
  STAT <- res.sum$coefficients[2, 1]/res.sum$coefficients[2, 2]
  table <- cbind(c(4.38, 4.15, 4.04, 3.99, 3.98, 3.96)
                 , c(3.95, 3.8, 3.73, 3.69, 3.68, 3.66)
                 , c(3.6, 3.5, 3.45, 3.43, 3.42, 3.41)
                 , c(3.24, 3.18, 3.15, 3.13, 3.13, 3.12)
                 , c(1.14, 1.19, 1.22, 1.23, 1.24, 1.25)
                 , c(0.8, 0.87, 0.9, 0.92, 0.93, 0.94)
                 , c(0.5, 0.58, 0.62, 0.64, 0.65, 0.66)
                 , c(0.15, 0.24, 0.28, 0.31, 0.32, 0.33))
  table <- -table
  tablen <- dim(table)[2]
  tableT <- c(25, 50, 100, 250, 500, 1e+05)
  tablep <- c(0.01, 0.025, 0.05, 0.1, 0.9, 0.95, 0.975, 0.99)
  tableipl <- numeric(tablen)
  for (i in (1:tablen)) tableipl[i] <- approx(tableT, table[, i], n, rule = 2)$y
  interpol <- approx(tableipl, tablep, STAT, rule = 2)$y
  if (alternative == "stationary") 
    PVAL <- round(interpol,3)
  else if (alternative == "explosive") 
    PVAL <- round(1 - interpol,3)
  else stop("irregular alternative")
  if (is.na(approx(tableipl, tablep, STAT, rule = 1)$y)) 
    if (interpol == min(tablep)) 
      PVAL <- paste("<",PVAL)
  else PVAL <- paste(">",PVAL)
  PARAMETER <- k - 1
  METHOD <- "Augmented Dickey-Fuller Test"
  names(STAT) <- "Dickey-Fuller"
  names(PARAMETER) <- "Lag order"
  structure(list(statistic = STAT, parameter = PARAMETER, alternative = alternative, 
                 p.value = PVAL, method = METHOD, data.name = DNAME), 
            class = "htest")
}


kpss.test2 <- function (x, null = c("Level", "Trend"), lshort = TRUE) 
{
  ## modification of output in kpss.test (include "<" and ">" in p-value output rather than warning)
  if ((NCOL(x) > 1) || is.data.frame(x)) 
    stop("x is not a vector or univariate time series")
  DNAME <- deparse(substitute(x))
  null <- match.arg(null)
  x <- as.vector(x, mode = "double")
  n <- length(x)
  if (null == "Trend") {
    t <- 1:n
    e <- residuals(lm(x ~ t))
    table <- c(0.216, 0.176, 0.146, 0.119)
  } else if (null == "Level") {
    e <- residuals(lm(x ~ 1))
    table <- c(0.739, 0.574, 0.463, 0.347)
  }
  tablep <- c(0.01, 0.025, 0.05, 0.1)
  s <- cumsum(e)
  eta <- sum(s^2)/(n^2)
  s2 <- sum(e^2)/n
  if (lshort) {
    l <- trunc(3 * sqrt(n)/13)
  } else l <- trunc(10 * sqrt(n)/14)
  s2 <- .C("tseries_pp_sum", as.vector(e, mode = "double"), as.integer(n), 
           as.integer(l), s2 = as.double(s2))$s2
  STAT <- eta/s2
  PVAL <- round(approx(table, tablep, STAT, rule = 2)$y,3)
  if (is.na(approx(table, tablep, STAT, rule = 1)$y)) 
    if (PVAL == min(tablep)) 
      PVAL <- paste("<",PVAL)
  else PVAL <- paste(">",PVAL)
  PARAMETER <- l
  METHOD <- paste("KPSS Test for", null, "Stationarity")
  names(STAT) <- paste("KPSS", null)
  names(PARAMETER) <- "Truncation lag parameter"
  structure(list(statistic = STAT, parameter = PARAMETER, p.value = PVAL, 
                 method = METHOD, data.name = DNAME), class = "htest")
}


stationarityTab <- function(x)
{
  ## Creates summary table of Dickey-Fuller test and KPSS test for multiple time series
  ## Arguments:
  ## - x: list of time series (names will be used as labels in the table)
  if(class(x) !="list") x <- list(x)
  tab <- NULL
  for(i in 1:length(x)){
    if(is.null(names(x)[i])) names(x)[i] <- paste("Variable",i)
    tab <- rbind(tab
                 , c(names(x)[i], "Dickey-Fuller", "Unit Root"
                     , adf.test2(na.omit(x[[i]]))$parameter
                     , round(adf.test2(na.omit(x[[i]]))$statistic,3)
                     , adf.test2(na.omit(x[[i]]))$p.value)
                 , c(NA, "KPSS", "Stationary"
                     , kpss.test2(na.omit(x[[i]]))$parameter
                     , round(kpss.test2(na.omit(x[[i]]))$statistic,3)
                     , kpss.test2(na.omit(x[[i]]))$p.value))
  }
  colnames(tab) <- c("Series","Test","$H_0$","Lags","Statistic","$p$-Value")
  return(tab)
}


## compute Ericsson/MacKinnon critical values based on response surface estimates and t
mkFunc <- function(k, t){
  ## k: number of variables in ecm (includes dv!)
  ## t: length of time series
  
  ## response surface estimates from Ericsson/MacKinnon (2002)
  ## (5% level with constant and no trend (Table 3))
  tab3 <- tribble(
    ~k, ~theta_inf, ~theta_1, ~theta_2, ~theta_3,
    1,   -2.8617,   -2.81,    -3.2,      37,
    2,   -3.2145,   -3.21,    -2.0,      17,
    3,   -3.5057,   -3.27,     1.1,     -34,
    4,   -3.7592,   -2.92,    -3.7,       5,
    5,   -3.9856,   -2.50,    -1.7,     -35,
    6,   -4.1922,   -1.73,    -7.8,      -9,
    7,   -4.3831,   -0.90,   -12.2,       1,
    8,   -4.5608,    0.02,   -15.4,      -2,
    9,   -4.7287,    1.25,   -26.0,      42,
    10,   -4.8876,    2.46,   -31.7,      43,
    11,   -5.0394,    3.88,   -45.7,     117,
    12,   -5.1836,    5.33,   -55.9,     134
  )
  
  ## compute adjusted sampe size (assuming constant w/o trend)
  ## (c.f., p.295)
  t_adj <- t - (2*k-1) - 1
  
  ## compute critical value
  crit <- tab3$theta_inf[k] + tab3$theta_1[k]/t_adj + 
    tab3$theta_2[k]/t_adj^2 + tab3$theta_3[k]/t_adj^3
  crit
}


## Wrapper for Amelia call
amelia_wrapper <- function(x, nimpute = 10, ktime = NULL,
                           lags = c("sqf_d", "ccrb_d", "arrests_d"),
                           leads = c("sqf_d", "ccrb_d", "arrests_d")){
  x %>% dplyr::select(date, sqf_d, ccrb_d, arrests_d, crime_d,  
                      media_d, unemp_d, visit_overseas_d, temp_mean, precip_total, jan) %>%
    amelia(ts = "date", polytime = ktime, m = nimpute, lags = lags, leads = leads)
}


## Function to print tables of VAR estimates
latexVAR <- function(x, caption=NULL, label=NULL, align=NULL, digits=3
                     , varlabs=NULL, mlabs=NULL, NeweyWest=FALSE, ...){
  ############################################################################################
  ### Function to print VAR model results in latex table
  ## x: model list
  ## caption: see xtable for details
  ## label: see xtable for details
  ## align: see xtable for details
  ## digits: nteger indicating the number of decimal places to be used
  ## varlabs: list of variable names and respective labels (also determines order in table)
  ## mlabs: list of variable names for each DV
  ## ...: further arguments passed to print.xtable (e.g. file etc.)
  ############################################################################################
  
  ## save model in list
  tbl <- data.frame(vars = colnames(x$datamat)[!colnames(x$datamat) %in% colnames(x$y)])
  
  ## arrange DVs according to mlab
  if(!is.null(mlabs)){
    if(length(mlabs)!=length(coef(x))) stop("Model labels do not have correct length")
    if(sum(!names(coef(x)) %in% names(mlabs))>0) stop("Model labels do not match DV names")
    dvnames <- names(mlabs)
  } else {
    dvnames <- names(coef(x))
  }
  
  for(dv in dvnames){
    ## extract varnames, coefs and se
    vars <- rownames(coef(x)[[dv]])
    coefs <- round(coef(x)[[dv]][,"Estimate"], digits)
    se <- paste0("(",round(coef(x)[[dv]][,"Std. Error"], digits),")")
    tmp <- data.frame(vars,coefs,se)
    colnames(tmp) <- c("vars",paste0(c("coefs_","se_"),dv))
    
    ## merge model results
    tbl <- merge(tbl,tmp,by="vars",all=T)
  }
  
  ## arrange variables
  if(!is.null(varlabs)){
    if(nrow(tbl)!=length(varlabs)) stop("Variable lables do not have correct length")
    rownames(tbl) <- as.character(tbl$vars)
    tbl <- tbl[names(varlabs),]
    tbl$vars <- unlist(varlabs)
  }
  
  ## prepare full variable names for output
  out <- data.frame(vars = as.vector(t(cbind(as.character(tbl$vars),""))))
  
  ## coefs and se in single column for each model
  for(m in dvnames){
    out <- cbind(out, as.vector(t(tbl[,grep(m,colnames(tbl))])))
  }
  
  ## model names
  if(!is.null(mlabs)){
    colnames(out) <- c("Variable", mlabs)
  } else {
    colnames(out) <- c("Variable", dvnames)
  }
  
  ## convert table to character
  out <- apply(out, 2, as.character)
  
  ## add number of observations
  out <- rbind(out, c("N",rep(x$obs, length(coef(x)))))
  
  ## add R-Squared
  rss <- apply(resid(x),2, var)
  tss <- apply(x$y[-(1:x$p),], 2, var)
  out <- rbind(out,  c("R-squared", round(1 - rss/tss, digits)))
  
  ## adjust align for excluded rownames
  if(!is.null(align)) align <- paste0("l",align)
  
  ## export table
  print(xtable(out, caption=caption, label=label, align=align), include.rownames=F
        , hline.after=c(-1,0,nrow(out)-2,nrow(out)), ...)
}

